#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Author  ：chenggl
@Date    ：2024/8/30 9:01
@DESC     ：baseline,采用词频IF/IDF进行特征提取，然后LR进行多分类
'''
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
import pandas as pd

label2id = {
    '安全隐患防范纠纷':0,
    '房屋宅基地纠纷':1,
    '环境保护纠纷':2,
    '交通事故纠纷':3,
    '教育培训纠纷':4,
    '经济合同纠纷':5,
    '劳动人事纠纷':6,
    '邻里纠纷':7,
    '其他纠纷':8,
    '侵权纠纷':9,
    '食品安全纠纷':10,
    '土地森林草场纠纷':11,
    '物业纠纷':12,
    '医疗纠纷':13,
    '噪音扰民纠纷':14,
    '征地拆迁纠纷':15
}

id2label = {
    0:'安全隐患防范纠纷',
    1:'房屋宅基地纠纷',
    2:'环境保护纠纷',
    3:'交通事故纠纷',
    4:'教育培训纠纷',
    5:'经济合同纠纷',
    6:'劳动人事纠纷',
    7:'邻里纠纷',
    8:'其他纠纷',
    9:'侵权纠纷',
    10:'食品安全纠纷',
    11:'土地森林草场纠纷',
    12:'物业纠纷',
    13:'医疗纠纷',
    14:'噪音扰民纠纷',
    15:'征地拆迁纠纷',
}

PATH_TRAIN = 'train.csv'
PATH_TEST = 'test.csv'

df_train = pd.read_csv(PATH_TRAIN,header=None)
print(df_train.head())
contents = df_train[1][1:]
labels = df_train[2][1:].apply(lambda x:label2id[x])

pipeline = make_pipeline(TfidfVectorizer(ngram_range=(2,2)),LogisticRegression())

pipeline.fit(contents,labels)
df_test = pd.read_csv(PATH_TEST,header=None)

content_tests = df_test[1][1:]
preds = pipeline.predict(content_tests)
pd.DataFrame({'ID':df_test[0][1:],'LABEL':preds}).to_csv('preds.csv', index=None)

